import numpy as np

# 1. 输入：最终向量（词嵌入+位置编码）
X = np.array([
    [1.0000, 1.0000, 0.0000, 1.0000],
    [0.8415, 1.5403, 0.0001, 1.0000],
    [0.9093, -0.4161, 1.0002, 0.9998]
])
print("输入最终向量 X:\n", np.round(X, 4), "\n")

# 2. 定义线性变换矩阵（Q、K、V），这里使用的是随机值，实际应用中会根据模型参数初始化，并在训练中不断更新。
W_Q = np.array([
    [0.2, 0.3, 0.1, 0.4],
    [0.5, 0.1, 0.2, 0.3],
    [0.3, 0.4, 0.5, 0.1],
    [0.1, 0.2, 0.3, 0.4]
])

W_K = np.array([
    [0.4, 0.1, 0.3, 0.2],
    [0.2, 0.5, 0.1, 0.3],
    [0.1, 0.3, 0.4, 0.5],
    [0.3, 0.2, 0.1, 0.4]
])

W_V = np.array([
    [0.1, 0.4, 0.2, 0.3],
    [0.3, 0.1, 0.5, 0.2],
    [0.2, 0.3, 0.1, 0.5],
    [0.4, 0.2, 0.3, 0.1]
])
# 2.1 定义输出投影矩阵，是将经过注意力计算后的向量映射到最终输出维度的线性变换矩阵
W_O = np.array([
    [0.3, 0.1, 0.2, 0.4],
    [0.1, 0.4, 0.3, 0.2],
    [0.2, 0.3, 0.4, 0.1],
    [0.4, 0.2, 0.1, 0.3]
])


# 3. 生成Q、K、V
Q = X @ W_Q
K = X @ W_K
V = X @ W_V
print("Q矩阵:\n", np.round(Q, 4), "\n")
print("K矩阵:\n", np.round(K, 4), "\n")
print("V矩阵:\n", np.round(V, 4), "\n")


# 4. 拆分多头（2头，每头维度2）
h = 2  # 头数
d_k = 4 // h  # 每个头的维度

Q_split = Q.reshape(3, h, d_k).transpose(1, 0, 2)  # 形状 (2,3,2)
K_split = K.reshape(3, h, d_k).transpose(1, 0, 2)
V_split = V.reshape(3, h, d_k).transpose(1, 0, 2)

print("头1 Q1:\n", np.round(Q_split[0], 4))
print("头1 K1:\n", np.round(K_split[0], 4))
print("头1 V1:\n", np.round(V_split[0], 4), "\n")

print("头2 Q2:\n", np.round(Q_split[1], 4))
print("头2 K2:\n", np.round(K_split[1], 4))
print("头2 V2:\n", np.round(V_split[1], 4), "\n")


# 5. 计算每个头的注意力
def scaled_dot_product_attention(Q, K, V):
    """
    实现缩放点积注意力机制
    参数:
        Q: 查询(Query)矩阵，形状通常为 (..., seq_len_q, d_k)
        K: 键(Key)矩阵，形状通常为 (..., seq_len_k, d_k)
        V: 值(Value)矩阵，形状通常为 (..., seq_len_v, d_v)，一般seq_len_k = seq_len_v
    返回:
        output: 注意力计算的输出，形状为 (..., seq_len_q, d_v)
        attn_weights: 注意力权重矩阵，形状为 (..., seq_len_q, seq_len_k)
    """
    # 获取查询向量的维度d_k（最后一个维度的大小）
    d_k = Q.shape[-1]
    # 1. 计算Q与K的转置的矩阵乘积，得到原始注意力分数；K.T表示对K的最后两个维度进行转置，使矩阵乘法合法
    scores = np.matmul(Q, K.T)  # 结果形状: (..., seq_len_q, seq_len_k)
    # 2. 对原始分数进行缩放，除以d_k的平方根，目的是缓解当d_k较大时，内积结果过大导致softmax梯度消失的问题
    scores = scores / np.sqrt(d_k)
    # 3. 对缩放后的分数应用softmax函数，得到归一化的注意力权重；axis=-1表示沿着最后一个维度（seq_len_k）进行归一化；keepdims=True保持维度，确保广播机制能正常工作
    attn_weights = np.exp(scores) / np.sum(np.exp(scores), axis=-1, keepdims=True)
    # 4. 将注意力权重与V进行矩阵乘法，得到最终的注意力输出，这一步是用注意力权重对值向量进行加权求和
    output = np.matmul(attn_weights, V)  # 结果形状: (..., seq_len_q, d_v)
    return output, attn_weights

# 头1计算
head1_output, attn1 = scaled_dot_product_attention(Q_split[0], K_split[0], V_split[0])
# 头2计算
head2_output, attn2 = scaled_dot_product_attention(Q_split[1], K_split[1], V_split[1])

print("头1注意力权重:\n", np.round(attn1, 4))
print("头1输出:\n", np.round(head1_output, 4), "\n")

print("头2注意力权重:\n", np.round(attn2, 4))
print("头2输出:\n", np.round(head2_output, 4), "\n")

# 6. 拼接多头输出
concatenated = np.concatenate([head1_output, head2_output], axis=1)
print("拼接后输出:\n", np.round(concatenated, 4), "\n")

# 7. 最终线性投影
multi_head_output = concatenated @ W_O
print("多头自注意力最终输出:\n", np.round(multi_head_output, 4))
    